"""
Modbus驱动实现
支持Modbus TCP和RTU协议
"""

import time
from datetime import datetime
from typing import List, Any, Dict

# pymodbus 动态导入 - 延迟到实际使用时
ModbusTcpClient = None
ModbusRtuClient = None
ModbusException = None
Endian = None
BinaryPayloadDecoder = None
BinaryPayloadBuilder = None

def _import_pymodbus():
    """动态导入pymodbus模块"""
    global ModbusTcpClient, ModbusRtuClient, ModbusException, Endian, BinaryPayloadDecoder, BinaryPayloadBuilder
    
    try:
        from pymodbus.client import ModbusTcpClient, ModbusRtuClient
        from pymodbus.exceptions import ModbusException
        from pymodbus.constants import Endian
        from pymodbus.payload import BinaryPayloadDecoder, BinaryPayloadBuilder
        return True
    except ImportError:
        try:
            # 尝试旧版本导入
            from pymodbus.client.tcp import ModbusTcpClient
            from pymodbus.client.rtu import ModbusRtuClient
            from pymodbus.exceptions import ModbusException
            from pymodbus.constants import Endian
            from pymodbus.payload import BinaryPayloadDecoder, BinaryPayloadBuilder
            return True
        except ImportError:
            return False

from ..base_driver import BaseDriver, TagValue, ConnectionStatus, DataType

class ModbusDriver(BaseDriver):
    """Modbus驱动类"""
    
    def __init__(self, config: Dict[str, Any]):
        super().__init__(config)
        
        # 在实际使用时再检查pymodbus可用性
        # if not PYMODBUS_AVAILABLE:
        #     raise ImportError("pymodbus库未安装，请运行: pip install pymodbus")
        
        self.client = None
        self.protocol = config.get('config', {}).get('protocol', 'tcp').lower()
        
        # TCP配置
        self.ip = config.get('config', {}).get('ip', 'localhost')
        self.port = config.get('config', {}).get('port', 502)
        
        # RTU配置
        self.serial_port = config.get('config', {}).get('serial_port', 'COM1')
        self.baudrate = config.get('config', {}).get('baudrate', 9600)
        self.parity = config.get('config', {}).get('parity', 'N')
        self.stopbits = config.get('config', {}).get('stopbits', 1)
        self.bytesize = config.get('config', {}).get('bytesize', 8)
        
        # 通用配置
        self.unit_id = config.get('config', {}).get('unit_id', 1)
        self.timeout = config.get('config', {}).get('timeout', 3)
        self.retries = config.get('config', {}).get('retries', 3)
        
        # 字节序配置
        self.word_order = config.get('config', {}).get('word_order', 'big')  # big, little
        self.byte_order = config.get('config', {}).get('byte_order', 'big')  # big, little
        
    def connect(self) -> bool:
        """连接Modbus设备"""
        try:
            # 动态导入pymodbus
            if not _import_pymodbus():
                self.logger.error("pymodbus库未安装，请运行: pip install pymodbus")
                return False
            if self.protocol == 'tcp':
                self.client = ModbusTcpClient(
                    host=self.ip,
                    port=self.port,
                    timeout=self.timeout,
                    retries=self.retries
                )
            elif self.protocol == 'rtu':
                self.client = ModbusRtuClient(
                    port=self.serial_port,
                    baudrate=self.baudrate,
                    parity=self.parity,
                    stopbits=self.stopbits,
                    bytesize=self.bytesize,
                    timeout=self.timeout,
                    retries=self.retries
                )
            else:
                self.logger.error(f"不支持的Modbus协议: {self.protocol}")
                return False
            
            # 尝试连接
            connection_result = self.client.connect()
            if connection_result:
                self.connection_status = ConnectionStatus.CONNECTED
                self.last_connection_time = datetime.now()
                self.last_error = None
                self.logger.info(f"Modbus {self.protocol.upper()} 连接成功")
                return True
            else:
                self.connection_status = ConnectionStatus.ERROR
                self.last_error = "连接失败"
                self.logger.error("Modbus连接失败")
                return False
                
        except Exception as e:
            self.connection_status = ConnectionStatus.ERROR
            self.last_error = str(e)
            self.logger.error(f"Modbus连接异常: {e}")
            return False
    
    def disconnect(self):
        """断开连接"""
        try:
            if self.client:
                self.client.close()
            self.connection_status = ConnectionStatus.DISCONNECTED
            self.logger.info("Modbus连接已断开")
        except Exception as e:
            self.logger.error(f"断开连接时出错: {e}")
    
    def read_tags(self, tag_names: List[str]) -> List[TagValue]:
        """批量读取标签值"""
        if not self.is_connected():
            return []
        
        tag_values = []
        
        for tag_name in tag_names:
            if tag_name not in self.data_points:
                self.logger.warning(f"未找到标签配置: {tag_name}")
                continue
                
            data_point = self.data_points[tag_name]
            
            try:
                raw_value = self._read_single_register(data_point.address, data_point.data_type)
                if raw_value is not None:
                    tag_value = self.process_tag_value(tag_name, raw_value)
                    tag_values.append(tag_value)
                else:
                    # 读取失败，创建错误标签值
                    tag_values.append(TagValue(
                        tag_name=tag_name,
                        value=None,
                        timestamp=datetime.now(),
                        quality="BAD",
                        device_name=self.device_name
                    ))
                    
            except Exception as e:
                self.logger.error(f"读取标签 {tag_name} 失败: {e}")
                tag_values.append(TagValue(
                    tag_name=tag_name,
                    value=None,
                    timestamp=datetime.now(),
                    quality="BAD",
                    device_name=self.device_name
                ))
        
        return tag_values
    
    def write_tag(self, tag_name: str, value: Any) -> bool:
        """写入标签值"""
        if not self.is_connected():
            return False
            
        if tag_name not in self.data_points:
            self.logger.error(f"未找到标签配置: {tag_name}")
            return False
            
        data_point = self.data_points[tag_name]
        
        # 检查写权限
        if 'W' not in data_point.access:
            self.logger.error(f"标签 {tag_name} 不允许写入")
            return False
        
        try:
            return self._write_single_register(data_point.address, value, data_point.data_type)
        except Exception as e:
            self.logger.error(f"写入标签 {tag_name} 失败: {e}")
            return False
    
    def _parse_address(self, address: str) -> tuple:
        """解析Modbus地址"""
        # 支持格式: 40001, 400001, 4:1, CO:1, HR:1 等
        if ':' in address:
            # 格式: HR:1, CO:1等
            func_code, addr = address.split(':')
            addr = int(addr) - 1  # Modbus地址从0开始
            
            if func_code.upper() in ['HR', 'HOLDING']:
                return 3, addr  # 读保持寄存器
            elif func_code.upper() in ['IR', 'INPUT']:
                return 4, addr  # 读输入寄存器
            elif func_code.upper() in ['CO', 'COIL']:
                return 1, addr  # 读线圈
            elif func_code.upper() in ['DI', 'DISCRETE']:
                return 2, addr  # 读离散输入
            else:
                raise ValueError(f"不支持的功能码: {func_code}")
        else:
            # 传统6位地址格式
            addr_int = int(address)
            if 400001 <= addr_int <= 465536:
                # 保持寄存器 (4x)
                return 3, addr_int - 400001
            elif 300001 <= addr_int <= 365536:
                # 输入寄存器 (3x)
                return 4, addr_int - 300001
            elif 100001 <= addr_int <= 165536:
                # 离散输入 (1x)
                return 2, addr_int - 100001
            elif 1 <= addr_int <= 65536:
                # 线圈 (0x)
                return 1, addr_int - 1
            elif 40001 <= addr_int <= 49999:
                # 保持寄存器简化格式
                return 3, addr_int - 40001
            elif 30001 <= addr_int <= 39999:
                # 输入寄存器简化格式
                return 4, addr_int - 30001
            else:
                raise ValueError(f"无效的Modbus地址: {address}")
    
    def _read_single_register(self, address: str, data_type: DataType) -> Any:
        """读取单个寄存器"""
        func_code, start_addr = self._parse_address(address)
        
        try:
            if func_code == 1:  # 读线圈
                result = self.client.read_coils(start_addr, 1, slave=self.unit_id)
                if result.isError():
                    self.logger.error(f"读取线圈失败: {result}")
                    return None
                return result.bits[0]
                
            elif func_code == 2:  # 读离散输入
                result = self.client.read_discrete_inputs(start_addr, 1, slave=self.unit_id)
                if result.isError():
                    self.logger.error(f"读取离散输入失败: {result}")
                    return None
                return result.bits[0]
                
            elif func_code in [3, 4]:  # 读寄存器
                # 根据数据类型确定寄存器数量
                register_count = self._get_register_count(data_type)
                
                if func_code == 3:
                    result = self.client.read_holding_registers(start_addr, register_count, slave=self.unit_id)
                else:
                    result = self.client.read_input_registers(start_addr, register_count, slave=self.unit_id)
                
                if result.isError():
                    self.logger.error(f"读取寄存器失败: {result}")
                    return None
                
                # 解码数据
                return self._decode_registers(result.registers, data_type)
                
        except ModbusException as e:
            self.logger.error(f"Modbus异常: {e}")
            return None
        except Exception as e:
            self.logger.error(f"读取寄存器异常: {e}")
            return None
    
    def _write_single_register(self, address: str, value: Any, data_type: DataType) -> bool:
        """写入单个寄存器"""
        func_code, start_addr = self._parse_address(address)
        
        try:
            if func_code == 1:  # 写线圈
                result = self.client.write_coil(start_addr, bool(value), slave=self.unit_id)
                return not result.isError()
                
            elif func_code == 3:  # 写保持寄存器
                registers = self._encode_value(value, data_type)
                if len(registers) == 1:
                    result = self.client.write_register(start_addr, registers[0], slave=self.unit_id)
                else:
                    result = self.client.write_registers(start_addr, registers, slave=self.unit_id)
                return not result.isError()
            else:
                self.logger.error(f"功能码 {func_code} 不支持写操作")
                return False
                
        except Exception as e:
            self.logger.error(f"写入寄存器异常: {e}")
            return False
    
    def _get_register_count(self, data_type: DataType) -> int:
        """获取数据类型对应的寄存器数量"""
        if data_type in [DataType.INT16, DataType.BOOL]:
            return 1
        elif data_type in [DataType.INT32, DataType.FLOAT]:
            return 2
        elif data_type == DataType.DOUBLE:
            return 4
        else:
            return 1
    
    def _decode_registers(self, registers: List[int], data_type: DataType) -> Any:
        """解码寄存器数据"""
        # 确定字节序
        word_endian = Endian.BIG if self.word_order == 'big' else Endian.LITTLE
        byte_endian = Endian.BIG if self.byte_order == 'big' else Endian.LITTLE
        
        decoder = BinaryPayloadDecoder.fromRegisters(
            registers, 
            byteorder=byte_endian, 
            wordorder=word_endian
        )
        
        if data_type == DataType.BOOL:
            return bool(registers[0])
        elif data_type == DataType.INT16:
            return decoder.decode_16bit_int()
        elif data_type == DataType.INT32:
            return decoder.decode_32bit_int()
        elif data_type == DataType.FLOAT:
            return decoder.decode_32bit_float()
        elif data_type == DataType.DOUBLE:
            return decoder.decode_64bit_float()
        elif data_type == DataType.STRING:
            return decoder.decode_string(len(registers) * 2).decode('utf-8').rstrip('\x00')
        else:
            return registers[0]
    
    def _encode_value(self, value: Any, data_type: DataType) -> List[int]:
        """编码值为寄存器数据"""
        word_endian = Endian.BIG if self.word_order == 'big' else Endian.LITTLE
        byte_endian = Endian.BIG if self.byte_order == 'big' else Endian.LITTLE
        
        builder = BinaryPayloadBuilder(byteorder=byte_endian, wordorder=word_endian)
        
        if data_type == DataType.BOOL:
            builder.add_16bit_int(1 if value else 0)
        elif data_type == DataType.INT16:
            builder.add_16bit_int(int(value))
        elif data_type == DataType.INT32:
            builder.add_32bit_int(int(value))
        elif data_type == DataType.FLOAT:
            builder.add_32bit_float(float(value))
        elif data_type == DataType.DOUBLE:
            builder.add_64bit_float(float(value))
        elif data_type == DataType.STRING:
            builder.add_string(str(value))
        else:
            builder.add_16bit_int(int(value))
        
        return builder.to_registers() 